import QUANTAXIS as QA

import os
import tushare as ts
import pandas as pd
import numpy as np
import time
import datetime


def avg(row):
    count = 0
    sum = 0
    for item in row:
        if item != -1 and item != 1:
            sum += item
            count += 1
    return sum / count


class AnalyzeReport(object):
    def __init__(self, codes, name='默认', start=None, end=None, calc_return=False):
        self.codes = codes
        self.name = name
        if start is None:
            self.start = '1996-01-01'
        else:
            self.start = start
        if end is None:
            self.end = str(datetime.date.today())
        else:
            self.end = end
        self._fetch_financial()
        self._fetch_stock_day(calc_return=calc_return)
        self.df_columns = [
                # 'date',
                # 'code',
                # 总市值 = 当前股价 * 总股本（238）
                'EP',       # 估值 EP 净利润（TTM）（276）/总市值 1
                'EPcut',    # 估值 EPcut 扣除非经常性损益后净利润（TTM）（206）/总市值 1
                'BP',       # 估值 BP 净资产/总市值 1 也等于 每股股价 / 每股净资产（004）
                'SP',       # 估值 SP 营业收入（TTM）（074 ? 283）/总市值 1
                'NCFP',     # 估值 NCFP 净现金流（TTM）/总市值 1 =
                'OCFP',     # 估值 OCFP 经营性现金流（TTM）/总市值 1 = 219每股经营性现金流 / 总股本（238）
                # 'DP',       # 估值 DP 近 12 个月现金红利（按除息日计）/总市值 1
                'G_PE',     # 估值 G/PE 净利润（TTM）同比增长率(184)/PE_TTM 1
                'Sales_G_q',   # 成长 Sales_G_q 营业收入（最新财报，YTD）同比增长率(183) 1
                'Profit_G_q',  # 成长 Profit_G_q 净利润（最新财报，YTD）同比增长率(184) 1
                # 'OCF_G_q',     # 成长 OCF_G_q 经营性现金流（最新财报，YTD）同比增长率 1
                               # = （本期经营活动现金净流量(225)-上期经营活动现金净流量）/上期经营活动现金净流量*100%
                'ROE_G_q',     # 成长 ROE_G_q ROE（最新财报，YTD）(006与上期相比)同比增长率 1
                'ROE_q',       # 财务质量 ROE_q ROE（最新财报，YTD）(006) 1
                'ROE_ttm',     # 财务质量 ROE_ttm ROE（最新财报，TTM）（006） 1
                'ROA_q',       # 财务质量 ROA_q ROA（最新财报，YTD）1 资产回报率= 净利润（276）/总资产（040）
                'ROA_ttm',     # 财务质量 ROA_ttm ROA（最新财报，TTM） 1
                'grossprofitmargin_q',       # 财务质量 grossprofitmargin_q 毛利率（最新财报，YTD）（202） 1
                # 'grossprofitmargin_ttm',     # 财务质量 grossprofitmargin_ttm 毛利率（最新财报，TTM） 1
                'profitmargin_q',            # 财务质量 profitmargin_q 扣除非经常性损益后净利润率（最新财报，YTD）（206） 1
                'profitmargin_ttm',          # 财务质量 profitmargin_ttm 扣除非经常性损益后净利润率（最新财报，TTM） 1
                'assetturnover_q',           # 财务质量 assetturnover_q 资产周转率（最新财报，YTD）（175） 1
                # 'assetturnover_ttm',         # 财务质量 assetturnover_ttm 资产周转率（最新财报，TTM） 1
                'operationcashflowratio_q',  # 财务质量 operationcashflowratio_q 经营性现金流/净利润（最新财报，YTD）（228） 1
                'operationcashflowratio_ttm',  # 财务质量 operationcashflowratio_ttm 经营性现金流/净利润（最新财报，TTM） 1
                'financial_leverage',    # 杠杆 financial_leverage 总资产（040）/净资产 -1 净资产 = 每股净资产（004） * 总股本（238）
                'debtequityratio',       # 杠杆 debtequityratio 非流动负债（062）/净资产 -1
                'cashratio',             # 杠杆 cashratio 现金比率（161） 161是错误的，使用公式 现金流/流动负债
                'currentratio',          # 杠杆 currentratio 流动比率（159） 159是错误的，使用公式 流动资产/流动负债
                'ln_capital'     # 市值 ln_capital 总市值取对数
                ]
        self.month_data = pd.date_range('1/1/1996', '12/31/2023', freq='M').astype(str).tolist()
        self.Q_data = pd.date_range('1/1/1996', '12/31/2023', freq='Q-MAR').astype(str).tolist()
        df_industry = pd.read_csv('D:\\PythonPro\\QUANTAXIS\\EXAMPLE\\AI\\data\\industry.csv', dtype=str)
        df_hy = df_industry[['code', 'name', 'industry']]
        df_hy = df_hy.set_index(['code'])
        self.codes = list(self.data_day.code)
        if len(self.codes) != df_hy.shape[0]:
            df_hy = df_hy.reindex(self.codes)
        self.data_industry = df_hy

    def __repr__(self):
        return '<Assess> %s' % self.name

    def _fetch_stock_day(self, calc_return=False):
        end_date = QA.QA_util_get_recent_months(self.end, months=2) if calc_return else self.end
        self.data_day = QA.QA_fetch_stock_day_adv(self.codes,
                                                  QA.QA_util_get_previous_year_month(self.start),
                                                  end_date)
        self.data_day = self.data_day.to_qfq()
        date_trade = self.end
        if not QA.QA_util_if_trade(date_trade):
            date_trade = QA.QA_util_get_pre_trade_date(date_trade)
        self.date_trade_datetime = pd.to_datetime(date_trade)
        self.stock_today = self.data_day.select_day(self.date_trade_datetime)
        self.data_hs300 = QA.QA_fetch_index_day_adv('000300',
                                                    self.start,
                                                    end_date)

    def _fetch_financial(self):
        self.data_financial = QA.QA_fetch_financial_report_adv(self.codes,
                                                               start=QA.QA_util_get_recent_months(QA.QA_util_get_previous_year(self.start), months=-6),
                                                               end=self.end)

    def score(self, date, df_bt=None):
        df = self.calc_financial_one_day(date)
        df = df[['code', 'score']]
        df = df.sort_values(by='score', ascending=False)
        print(df.head())
        print(df_bt.head())
        if df_bt is not None:
            df = df_bt.merge(df, how='inner', left_on='code', right_on='code')
        df = df[['code', 'end', 'score']].sort_values(by='score', ascending=False)
        print(df.head())
        df.to_csv('bt_score.csv')

    def calc_financial_one_day(self, date):
        print('Time:%s' % date)
        date_trade = QA.QA_util_get_real_date(date, towards=-1)

        date_financial = QA.QA_util_get_real_fina_date(date)
        print('The real financial date is %s' % date_financial)
        if date_financial in self.Q_data:
            # 计算同比季报
            date_financial_datetime = pd.to_datetime(date_financial)
            date_financial_previous = self.Q_data[self.Q_data.index(date_financial) - 4]
            date_financial_previous_datetime = pd.to_datetime(date_financial_previous)
            # 计算TTM
            ttm_list = QA.QA_util_get_TTM_date_list(date)
            # print(ttm_list)
            date_financial_ttm = [pd.to_datetime(x) for x in ttm_list]
        if date_financial_datetime is None:
            print('%s未计算，因为不在Q期内' % date)
            return

        date_trade_datetime = pd.to_datetime(date_trade)

        stock_today = self.data_day.select_day(date_trade_datetime)
        print('截面%s的代码共%d个' % (date_trade, len(codes)))  # 每个截面的代码个数是不一样的，将codes放在截面里，便于提取

        close = stock_today.CLOSE
        res_data = self.data_financial.data.loc[date_financial_datetime]
        res_data_pre = self.data_financial.data.loc[date_financial_previous_datetime]
        res_data_ttm = self.data_financial.data.loc[date_financial_ttm].drop(['code', 'report_date'], axis=1)
        res_data_ttm = res_data_ttm.reset_index().groupby('code').sum()

        totalCapital = res_data.totalCapital + 1  # 238 总股本
        totalValue = close * totalCapital
        netAssetsPerShare = res_data.netAssetsPerShare  # 每股净资产（004）
        netAssets = netAssetsPerShare * totalCapital  # 净资产
        ROE = res_data.netAssetsPerShare  # （006）
        totalAssets = res_data.totalAssets  # 040 总资产
        totalNonCurrentLiabilities = res_data.totalNonCurrentLiabilities  # 062
        netCashFlows = res_data.netCashFlowsFromOperatingActivities  # 107经营活动产生的现金流量净额
        netProfitGrowthRate = res_data.netProfitGrowthRate  # 184
        cutProfit = res_data.netProfitAfterExtraordinaryGainsAndLosses  # 206
        netCashFlowPerShare = res_data.netCashFlowPerShare  # 225
        operatingCashFlowPerShare = res_data.operatingCashFlowPerShare  # 219
        listedAShares = res_data.listedAShares  # 239 A股流通股本
        lastIncome = res_data.lastYearOperatingIncome  # 283
        netProfit = res_data.netProfitsBelongToParentCompanyOwner  # 096
        totalOwnersEquity = res_data.totalOwnersEquity  # 072 所有者权益
        operatingRevenue = res_data.operatingRevenue  # 074 营业收入
        totalLiquidAssets = res_data.totalLiquidAssets  # 021流动资产合计
        totalCurrentLiabilities = res_data.totalCurrentLiabilities  # 054流动负债合计
        netCashFlows = res_data.netCashFlowsFromOperatingActivities \
                       + res_data.netCashFlowsFromInvestingActivities \
                       + res_data.netCashFlowsFromFinancingActivities

        # TTM
        operatingRevenue_TTM = res_data_ttm.operatingRevenueSingle  # 230 营业收入
        operatingProfit_TTM = res_data_ttm.operatingProfitSingle  # 231 营业利润
        netProfit_TTM = res_data_ttm.netProfitBelongingToTheOwnerOfTheParentCompanySingle  # 232
        cutProfit_TTM = res_data_ttm.netProfitAfterExtraordinaryGainsAndLossesSingle  # 233
        netCashFlowsFromOperatingActivities_TTM = res_data_ttm.netCashFlowsFromOperatingActivitiesSingle  # 234
        netCashFlowsFromInvestingActivities_TTM = res_data_ttm.netCashFlowsFromInvestingActivitiesSingle  # 235
        netCashFlowsFromFinancingActivities_TTM = res_data_ttm.netCashFlowsFromFinancingActivitiesSingle  # 236
        netIncreaseInCashAndCashEquivalents_TTM = res_data_ttm.netIncreaseInCashAndCashEquivalentsSingle  # 237
        # TODO 净现金流目前使用三项相加，不知道对错
        netCashFlows_TTM = netCashFlowsFromOperatingActivities_TTM + netCashFlowsFromInvestingActivities_TTM + netCashFlowsFromFinancingActivities_TTM
        pre_ROE = res_data_pre.netAssetsPerShare  # （006）
        pre_netCashFlowPerShare = res_data_pre.netCashFlowPerShare  # 225
        print('已获取财务数据')

        df_factors = pd.DataFrame({
            # 'date': date,
            # 'code': data_d.code,
            # 'name': df_hy.name[codes],
            # 'industry': df_hy.c_name[codes],
            # 'label': label,
            # 'label_return': label_return,
            'EP': netProfit_TTM / totalValue,
            'EPcut': cutProfit_TTM / totalValue,
            'BP': close / netAssetsPerShare,
            'SP': operatingRevenue_TTM / totalValue,
            'NCFP': netCashFlows_TTM / totalValue,
            'OCFP': netCashFlowsFromOperatingActivities_TTM / totalCapital,
            # 'DP': 0,
            'G_PE': netProfitGrowthRate * totalValue / netProfit,
            'Sales_G_q': self.data_financial.data.operatingIncomeGrowth[date_financial_datetime],
            'Profit_G_q': netProfitGrowthRate,
            # 'OCF_G_q': (netCashFlowPerShare - pre_netCashFlowPerShare) / pre_netCashFlowPerShare,  # 可能除0
            'ROE_G_q': (ROE - pre_ROE) / pre_ROE,
            'ROE_q': ROE,
            'ROE_ttm': netProfit_TTM / totalOwnersEquity,
            'ROA_q': netProfit / totalAssets,
            'ROA_ttm': netProfit_TTM / totalAssets,
            'grossprofitmargin_q': self.data_financial.data.rateOfReturnOnGrossProfitFromSales[date_financial_datetime],
            # 'grossprofitmargin_ttm': 0,
            'profitmargin_q': cutProfit / operatingRevenue,
            'profitmargin_ttm': cutProfit_TTM / operatingRevenue_TTM,
            'assetturnover_q': self.data_financial.data.turnoverRatioOfTotalAssets[date_financial_datetime],
            # 'assetturnover_ttm': 0,
            'operationcashflowratio_q': self.data_financial.data.cashFlowRateAndNetProfitRatioOfOperatingActivities[
                date_financial_datetime],
            'operationcashflowratio_ttm': netCashFlowsFromOperatingActivities_TTM / netProfit_TTM,
            'financial_leverage': netAssets / totalAssets,  # 为便于计算分数，这里取倒数
            'debtequityratio': netAssets / totalNonCurrentLiabilities,  # 为便于计算分数，这里取倒数
            # 'cashratio': self.data_financial.data.cashRatio[date_financial_datetime],
            'cashratio': netCashFlows / totalCurrentLiabilities,
            # 'currentratio': self.data_financial.data.currentRatio[date_financial_datetime],
            'currentratio': totalLiquidAssets / totalCurrentLiabilities,
            'ln_capital': np.log(totalValue)
        }, index=codes)
        df_tmp = pd.DataFrame(index=codes)
        df_factors = df_factors.replace([np.inf, -np.inf], np.nan)

        # 去极值 MAD法
        df_factors = df_factors.apply(QA.QA_util_median_limit)

        # 按行业进行均值填充nan
        df_factors['industry'] = self.data_industry.industry[codes]
        df_factors.to_csv('test.csv')

        df_factors = df_factors.groupby('industry').transform(
            lambda x: x.fillna(
                0 if np.isnan(x.mean()) else x.mean()  # 有些行业的因子全部都为空
            ))

        # 行业市值中性化
        df_factors = df_factors.apply(
            lambda x: x if x.name == 'code' else QA.QA_util_neutralization(x,
                                                                           df_hy=self.data_industry.loc[codes],
                                                                           # total_value=None))
                                                                           total_value=df_factors['ln_capital']))

        # z-score标准化
        df_factors = df_factors.apply(lambda x: x if x.name == 'code' else (x - x.mean()) / x.std())

        # 裁剪数值到-1~1
        df_factors = df_factors.clip(lower=-1, upper=1)

        # 主成分分析
        # reduced_x = pca.fit_transform(df_tmp[df_columns])
        # df_tmp[df_columns] = reduced_x

        # 计算分数, 简单求和
        df_factors['score'] = df_factors.apply(lambda x: avg(x), axis=1)

        df_tmp['date'] = date
        df_tmp['name'] = self.data_industry.name[codes]
        df_tmp['close'] = close
        df_tmp['industry'] = self.data_industry.industry[codes]
        df_tmp.dropna(inplace=True)
        df_tmp.index.name = 'code'
        df_tmp = pd.concat([df_tmp, df_factors.reindex(df_tmp.index)], axis=1)
        df_tmp.reset_index(inplace=True)
        # df_tmp.set_index(['date', 'code'], inplace=True)
        df_tmp.to_csv('%s_fina.csv' % date, encoding='utf_8_sig')
        return df_tmp

    def calc_data(self):
        series = pd.Series(data=self.month_data, index=pd.to_datetime(self.month_data), name='date')
        timerange = series.loc[self.start: self.end].tolist()
        df = pd.DataFrame(columns=self.df_columns)

        date_financial_datetime = None
        for date in timerange:
            df_tmp = self.calc_financial_one_day(date)
            df = df.append(df_tmp, sort=False)

        self.data = df.set_index(['date', 'code'])
        print(self.data.head(10))
        self.data.to_csv('financial.csv', encoding='utf_8_sig')

    def calc_data_with_return(self):
        series = pd.Series(data=self.month_data, index=pd.to_datetime(self.month_data), name='date')
        timerange = series.loc[self.start: self.end].tolist()
        df = pd.DataFrame(columns=self.df_columns)
        df_class = pd.DataFrame()

        date_financial_datetime = None
        index = 0
        for date in timerange:
            print('Time:%s' % date)
            date_next_month_last = QA.QA_util_get_real_date(QA.QA_util_get_recent_months(date, 1), towards=-1)
            date_trade = QA.QA_util_get_real_date(date, towards=-1)
            date_next_month_first = QA.QA_util_get_real_date(QA.QA_util_get_month_day_1(date_next_month_last),
                                                             towards=1)
            date_next_next_month_first = QA.QA_util_get_real_date(
                QA.QA_util_get_month_day_1(QA.QA_util_get_recent_months(date, 2)), towards=1)

            date_financial = QA.QA_util_get_real_fina_date(date)
            print('The real financial date is %s' % date_financial)
            if date_financial in self.Q_data:
                # 计算同比季报
                date_financial_datetime = pd.to_datetime(date_financial)
                date_financial_previous = self.Q_data[self.Q_data.index(date_financial) - 4]
                date_financial_previous_datetime = pd.to_datetime(date_financial_previous)
                # 计算TTM
                ttm_list = QA.QA_util_get_TTM_date_list(date)
                # print(ttm_list)
                date_financial_ttm = [pd.to_datetime(x) for x in ttm_list]
            if date_financial_datetime is None:
                print('%s未计算，因为不在Q期内' % date)
                continue

            date_trade_datetime = pd.to_datetime(date_trade)
            date_trade_next_month_datetime = pd.to_datetime(date_next_month_last)
            date_trade_next_month_first_datetime = pd.to_datetime(date_next_month_first)
            date_trade_next_next_month_first_datetime = pd.to_datetime(date_next_next_month_first)

            stock_today = self.data_day.select_day(date_trade_datetime)
            stock_next_month_last = self.data_day.select_day(date_trade_next_month_datetime)
            stock_next_month_first = self.data_day.select_day(date_trade_next_month_first_datetime)
            stock_next_next_month_first = self.data_day.select_day(date_trade_next_next_month_first_datetime)
            codes = list(set(stock_today.data.index).intersection(list(stock_next_month_last.data.index),
                                                                  list(stock_next_month_first.data.index),
                                                                  list(stock_next_next_month_first.data.index)))
            print('截面%s的代码共%d个' % (date_trade, len(codes)))  # 每个截面的代码个数是不一样的，将codes放在截面里，便于提取

            index_today = self.data_hs300.select_day(date_trade_datetime)
            index_next_month = self.data_hs300.select_day(date_trade_next_month_datetime)
            index_return = (index_next_month.CLOSE - index_today.CLOSE) / index_today.CLOSE
            index_return = index_return.iloc[0]

            # stock_today = data_d.data.loc[date_trade_datetime]
            # stock_today['CLOSE'] = stock_today.close.fillna(stock_today.preclose)
            # stock_next_month = data_d.data.loc[date_trade_next_month_datetime]
            # stock_next_month['CLOSE'] = stock_next_month.close.fillna(stock_next_month.preclose)
            forward_return = (
                                         stock_next_next_month_first.CLOSE - stock_next_month_first.CLOSE) / stock_next_month_first.CLOSE
            label_return = (stock_next_month_last.CLOSE - stock_today.CLOSE) / stock_today.CLOSE
            label = label_return - index_return

            close = stock_today.CLOSE
            res_data = self.data_financial.data.loc[date_financial_datetime]
            res_data_pre = self.data_financial.data.loc[date_financial_previous_datetime]
            res_data_ttm = self.data_financial.data.loc[date_financial_ttm].drop(['code', 'report_date'], axis=1)
            res_data_ttm = res_data_ttm.reset_index().groupby('code').sum()

            totalCapital = res_data.totalCapital + 1  # 238 总股本
            totalValue = close * totalCapital
            netAssetsPerShare = res_data.netAssetsPerShare  # 每股净资产（004）
            netAssets = netAssetsPerShare * totalCapital  # 净资产
            ROE = res_data.netAssetsPerShare  # （006）
            totalAssets = res_data.totalAssets  # 040 总资产
            totalNonCurrentLiabilities = res_data.totalNonCurrentLiabilities  # 062
            netCashFlows = res_data.netCashFlowsFromOperatingActivities  # 107经营活动产生的现金流量净额
            netProfitGrowthRate = res_data.netProfitGrowthRate  # 184
            cutProfit = res_data.netProfitAfterExtraordinaryGainsAndLosses  # 206
            netCashFlowPerShare = res_data.netCashFlowPerShare  # 225
            operatingCashFlowPerShare = res_data.operatingCashFlowPerShare  # 219
            listedAShares = res_data.listedAShares  # 239 A股流通股本
            lastIncome = res_data.lastYearOperatingIncome  # 283
            netProfit = res_data.netProfitsBelongToParentCompanyOwner  # 096
            totalOwnersEquity = res_data.totalOwnersEquity  # 072 所有者权益
            operatingRevenue = res_data.operatingRevenue  # 074 营业收入

            # TTM
            operatingRevenue_TTM = res_data_ttm.operatingRevenueSingle  # 230 营业收入
            operatingProfit_TTM = res_data_ttm.operatingProfitSingle  # 231 营业利润
            netProfit_TTM = res_data_ttm.netProfitBelongingToTheOwnerOfTheParentCompanySingle  # 232
            cutProfit_TTM = res_data_ttm.netProfitAfterExtraordinaryGainsAndLossesSingle  # 233
            netCashFlowsFromOperatingActivities_TTM = res_data_ttm.netCashFlowsFromOperatingActivitiesSingle  # 234
            netCashFlowsFromInvestingActivities_TTM = res_data_ttm.netCashFlowsFromInvestingActivitiesSingle  # 235
            netCashFlowsFromFinancingActivities_TTM = res_data_ttm.netCashFlowsFromFinancingActivitiesSingle  # 236
            netIncreaseInCashAndCashEquivalents_TTM = res_data_ttm.netIncreaseInCashAndCashEquivalentsSingle  # 237
            # TODO 净现金流目前使用三项相加，不知道对错
            netCashFlows_TTM = netCashFlowsFromOperatingActivities_TTM + netCashFlowsFromInvestingActivities_TTM + netCashFlowsFromFinancingActivities_TTM
            pre_ROE = res_data_pre.netAssetsPerShare  # （006）
            pre_netCashFlowPerShare = res_data_pre.netCashFlowPerShare  # 225
            print('已获取财务数据')

            df_factors = pd.DataFrame({
                # 'date': date,
                # 'code': data_d.code,
                # 'name': df_hy.name[codes],
                # 'industry': df_hy.c_name[codes],
                # 'label': label,
                # 'label_return': label_return,
                'EP': netProfit_TTM / totalValue,
                'EPcut': cutProfit_TTM / totalValue,
                'BP': close / netAssetsPerShare,
                'SP': operatingRevenue_TTM / totalValue,
                'NCFP': netCashFlows_TTM / totalValue,
                'OCFP': netCashFlowsFromOperatingActivities_TTM / totalCapital,
                # 'DP': 0,
                'G_PE': netProfitGrowthRate * totalValue / netProfit,
                'Sales_G_q': self.data_financial.data.operatingIncomeGrowth[date_financial_datetime],
                'Profit_G_q': netProfitGrowthRate,
                # 'OCF_G_q': (netCashFlowPerShare - pre_netCashFlowPerShare) / pre_netCashFlowPerShare,  # 可能除0
                'ROE_G_q': (ROE - pre_ROE) / pre_ROE,
                'ROE_q': ROE,
                'ROE_ttm': netProfit_TTM / totalOwnersEquity,
                'ROA_q': netProfit / totalAssets,
                'ROA_ttm': netProfit_TTM / totalAssets,
                'grossprofitmargin_q': self.data_financial.data.rateOfReturnOnGrossProfitFromSales[date_financial_datetime],
                # 'grossprofitmargin_ttm': 0,
                'profitmargin_q': cutProfit / operatingRevenue,
                'profitmargin_ttm': cutProfit_TTM / operatingRevenue_TTM,
                'assetturnover_q': self.data_financial.data.turnoverRatioOfTotalAssets[date_financial_datetime],
                # 'assetturnover_ttm': 0,
                'operationcashflowratio_q': self.data_financial.data.cashFlowRateAndNetProfitRatioOfOperatingActivities[
                    date_financial_datetime],
                'operationcashflowratio_ttm': netCashFlowsFromOperatingActivities_TTM / netProfit_TTM,
                'financial_leverage': totalAssets / netAssets,
                'debtequityratio': totalNonCurrentLiabilities / netAssets,
                'cashratio': self.data_financial.data.cashRatio[date_financial_datetime],
                'currentratio': self.data_financial.data.currentRatio[date_financial_datetime],
                'ln_capital': np.log(totalValue)
                }, index=codes)
            df_tmp = pd.DataFrame(index=codes)

            # df_tmp = df_tmp.dropna(subset=['label'])  # 删除标签为空的行, 加入后添加code列时会无法对齐, 以后想办法处理
            # 去极值 MAD法
            df_factors = df_factors.apply(QA.QA_util_median_limit)
            # df_tmp['code'] = codes

            # 按行业进行均值填充nan
            df_factors['industry'] = self.data_industry.industry[codes]
            df_factors = df_factors.groupby('industry').transform(lambda x: x.fillna(x.mean()))

            # 行业市值中性化
            df_factors = df_factors.apply(
                lambda x: x if x.name == 'code' else QA.QA_util_neutralization(x, df_hy=self.data_industry.loc[codes],
                                                                               total_value=df_factors['ln_capital']))

            # z-score标准化
            df_factors = df_factors.apply(lambda x: x if x.name == 'code' else (x - x.mean()) / x.std())

            # 主成分分析
            # reduced_x = pca.fit_transform(df_tmp[df_columns])
            # df_tmp[df_columns] = reduced_x

            df_tmp['date'] = date
            df_tmp['name'] = self.data_industry.name[codes]
            df_tmp['close'] = close
            df_tmp['label'] = label
            df_tmp['label_return'] = label_return
            df_tmp['forward_return'] = forward_return
            df_tmp['industry'] = self.data_industry.industry[codes]
            df_tmp.dropna(inplace=True)
            df_tmp.index.name = 'code'
            df_cate = df_tmp.sort_values(by='label', ascending=False)
            df_cate = pd.concat([df_cate, df_factors.reindex(df_cate.index)], axis=1)
            df_cate.reset_index(inplace=True)
            samples_count = int(df_cate.shape[0] * 0.3)
            df_true = df_cate.iloc[:samples_count]
            df_false = df_cate.iloc[-samples_count:]
            df_true.loc[:, 'label_cate'] = 1
            df_false.loc[:, 'label_cate'] = 0
            df_tmp = pd.concat([df_tmp, df_factors.reindex(df_tmp.index)], axis=1)
            df_tmp.reset_index(inplace=True)
            df = df.append(df_tmp, sort=False)
            df_class = df_class.append(df_true, sort=False)
            df_class = df_class.append(df_false, sort=False)
            index += 1

        self.data = df.set_index(['date', 'code'])
        print(self.data.head(10))
        self.data.to_csv('train.csv', encoding='utf_8_sig')


if __name__ == '__main__':
    codes = ['000002', '000100']

    # 300+500
    file_300 = 'D:\\PythonPro\\QUANTAXIS\\EXAMPLE\\AI\\DATA\\hs300.csv'
    if os.path.exists(file_300):
        hs300 = pd.read_csv(file_300, dtype=str)
    else:
        hs300 = ts.get_hs300s()
    file_500 = 'D:\\PythonPro\\QUANTAXIS\\EXAMPLE\\AI\\DATA\\zz500.csv'
    if os.path.exists(file_500):
        zz500 = pd.read_csv(file_500, dtype=str)
    else:
        zz500 = ts.get_zz500s()
    codes = pd.concat([hs300, zz500])['code'].to_list()  # zz500 + hs300

    # 全部hs
    ts_token = '17056d23a59ab71cb979c6a30185e092aba605c4544dac900a3eb7f8'
    ts.set_token(ts_token)
    pro = ts.pro_api()
    data = pro.stock_basic(exchange='', list_status='L', fields='symbol,name,area,industry,list_date')
    data = data[data.symbol.str.startswith('00') | data.symbol.str.startswith('60') | data.symbol.str.startswith('30')]
    # data = data[data.symbol.str.startswith('600085')]
    data = data[data.list_date < '20160101']
    all = data.rename(columns={'symbol': 'code'})
    all = all[~all.name.str.contains('ST')]
    # codes = all['code'].to_list()
    # codes.remove('000029')

    start_date = '2020-03-01'
    end_date = str(datetime.date.today() + datetime.timedelta(days=-1))
    ar = AnalyzeReport(codes=codes, start=start_date, end=end_date)
    # df_bt = pd.read_csv('break_through.csv', dtype=str)
    # ar.score(end_date, df_bt=df_bt)
    ar.calc_financial_one_day('2021-02-23')
    # ar.calc_data()